#! /usr/bin/env python
#coding=utf-8

"""

    2011 Boris Tatarintsev

    Simple tree data structure, very handy to use with the big bulk of data 
    such as the structure of livejournal users

"""

class HashTree:
    
    def __init__(self):
        self.m_hashTable = {}
        self.m_maxDepth = -1
    
    def __str__(self):
        output = str(self.m_hashTable)
        return output
    
    def getKeys(self):
        return self.m_hashTable.keys()
    
    def setMaxDepth(self, depth):
        self.m_maxDepth = depth
        
    def getMaxDepth(self):
        return self.m_maxDepth
    
    def setRef(self, fromNode, toNode):
        if self.m_hashTable.has_key(fromNode):
            fromRef = self.m_hashTable[fromNode]
            fromRef.m_refSet.add(toNode)
        else:
            # return 0 if the given node doesn't exist 
            return 0

    def hasNode(self, nodeName):
        return self.m_hashTable.has_key(nodeName)
    
    def getRefSet(self, nodeName):
        if self.m_hashTable.has_key(nodeName):
            return self.m_hashTable[nodeName].m_refSet
        return None
    
    def getNodeData(self, nodeName, key):
        if self.m_hashTable.has_key(nodeName):
            if self.m_hashTable[nodeName].m_data.has_key(key):
                return self.m_hashTable[nodeName].m_data[key]
        return None

    def filterNodesByProperty(self, property, value):
        output = []
        for key in self.m_hashTable.keys():
            v = self.getNodeData(key, property)
            if v != None and v == value:
                output.append(self.m_hashTable[key])                
        return output

    def getPlainHash(self):
        return self.m_hashTable

    def addNode(self, nodeName):
        # add node to the initial hash of nodes
        if not self.m_hashTable.has_key(nodeName):
            self.m_hashTable[nodeName] = Node()
            self.m_hashTable[nodeName].m_name = nodeName
        return self.m_hashTable[nodeName]
            
    def addNodeData(self, nodeName, key, value):
        if self.m_hashTable.has_key(nodeName):
            self.m_hashTable[nodeName].m_data[key] = value
            return 1
        return 0
    
    def removeNodeData(self, nodeName, key):
        if self.getNodeData(nodeName, key) != None:        
            del self.m_hashTable[nodeName].m_data[key]
    
    def setDepth(self, nodeName, depth):        
        self.addNodeData(nodeName, 'D', depth)
    
    def getDepth(self, nodeName):
        return self.getNodeData(nodeName, 'D')
    
    def getNode(self, nodeName):
        # get the node associated with the given name
        if self.hasNode(nodeName):            
            return self.m_hashTable[nodeName]
        else: return None

    def removeNode(self, nodeName):
        # remove a node from the tree
        if self.m_hashTable.has_key(nodeName):
            del self.m_hashTable[nodeName]
            return 1        
        return 0
    
    def findRoot(self):
        rootNode = None        
        for key in self.m_hashTable.keys():
            if self.getNodeData(key, 'ROOT') == 'True':                
                rootNode = self.getNode(key)
                return rootNode
        return None

    def toHR(self):        
        return self.__traverse(self.findRoot(), set(), [])
                
    def __traverse(self, node, visited, output):
        visited.add(node)
        output.append(str(node))
        for ref in node.m_refSet:
            if self.getNode(ref) not in visited:
                self.__traverse(self.getNode(ref), visited, output)
        return output

class Node:
    
    def __init__(self):
        # use set for references container to avoid duplicates
        # and improve performance
        self.m_refSet = set()
        # user defined data
        self.m_data = {}
        # parent ref
        self.m_name = ""
        
    def __str__(self):
        # string representation of a node
        output = 'name: ' + str(self.m_name) + '\nreferences: ' + str(self.m_refSet) + '\n'
        return output


# test case
"""
tree = HashTree()

tree.addNode('user1')
tree.addNodeData('user1', 'ROOT', 'True')
tree.addNode('user2')
tree.addNode('user3')
tree.addNode('user4')

tree.setRef('user1', 'user3')
tree.setRef('user1', 'user2')
tree.setRef('user3', 'user1')
tree.setRef('user2', 'user4')
print tree.toString()
"""